#!/usr/bin/env python3

from random import randint
import argparse, sys

num_osd = 100
num_groups = 10
num_pg_members = 3
num_pg = 16384
num_failures = 1000000

pg_list = {}	# {pg0:[osd1, osd2, osd], pg1: [osd3, osd5, osd9], ...}
osd_list = {}	# {osd0:[pg0, pg8, pg7, ...], osd1:[pg2, pg4, ...], ...}
osd_grps = {}	# {grp0:[osd1, osd2, ...], grp3: [osd10, osd11, ...], ...}

def divide_osd_into_groups():
	num_osd_per_grp = num_osd // num_groups
	gid = 0
	for osd in range(num_osd):
		grp = osd_grps.setdefault(gid, [])
		grp.append(osd)
		if gid < num_groups-1 and len(grp) >= num_osd_per_grp:
			gid += 1
	return

def assign_osd_to_pgs():
	for pg_id in range(num_pg):
		pg_list.setdefault(pg_id, [])
		gid = randint(0, num_groups-1)
		while len(pg_list[pg_id]) < num_pg_members:
			idx = randint(0, len(osd_grps[gid])-1)
			osd = osd_grps[gid][idx]
			if osd in pg_list[pg_id]:
				continue
			pg_list[pg_id].append(osd)
			osd_list.setdefault(osd, [])
			osd_list[osd].append(pg_id)
	return

def sim_once_failure():
	'''
	return True if data lost happens
	'''
	fail_osd = []
	while len(fail_osd) < num_pg_members:
		osd = randint(0, num_osd-1)
		if osd in fail_osd:
			continue
		fail_osd.append(osd)

	#
	# hit_pg is the intersection of the 3 SETs of PGs which are hosted
	# by the 3 failed OSDs.
	#
	hit_pg = set(osd_list[fail_osd[0]])
	for i in range(1, num_pg_members):
		hit_pg = hit_pg.intersection(set(osd_list[fail_osd[i]]))
	#
	# if that hit_pg is not empty, then data lost happens
	#
	return hit_pg != set()
	

def main():
	global num_osd, num_groups, num_pg
	parser = argparse.ArgumentParser()
	parser.add_argument('-g', action='store', dest='num_groups', help='number of groups')
	parser.add_argument('-o', action='store', dest='num_osd', help='number of OSDs')
	parser.add_argument('-p', action='store', dest='num_pg', help='number of PGs')
	args = parser.parse_args()

	if args.num_osd:
		num_osd = int(args.num_osd)
		if num_osd<num_pg_members:
			print('number of osd must be greater than %d' % num_pg_members)
			sys.exit(0)
	
	max_groups = num_osd // num_pg_members
	if num_groups > max_groups:
		num_groups = max_groups

	if args.num_groups:
		num_groups = int(args.num_groups)
		if num_groups<=0 or num_groups>=max_groups:
			print('number of groups must be in the range [1, %d]'% max_groups)
			sys.exit(0)

	if args.num_pg:
		num_pg = int(args.num_pg)
		if num_pg<0 or num_pg>16384:
			print('number of PGs must be in the range [1, 16384]')
			sys.exit(0)
	
	print("num_osd = %d" % num_osd)
	print("num_groups = %d" % num_groups)
	print("num_pg_members = %d" % num_pg_members)
	print("num_pg = %d" % num_pg)
	print("num_failures = %d" % num_failures)

	divide_osd_into_groups()
	assign_osd_to_pgs()
	num_data_lost = 0
	for i in range(num_failures):
		if (i+1) % 1000 == 0:
			probability = num_data_lost / (i+1)
			print("failures = %d, data lost = %d, probability_data_lost = %f" % (i+1, num_data_lost, probability))
		data_lost = sim_once_failure()
		if data_lost:
			num_data_lost += 1
	probability = num_data_lost / num_failures
	print("probability_data_lost = %f" % probability)

if __name__=="__main__":
	try:
		main()
	except KeyboardInterrupt:
		print('Exit')
		sys.exit(0)

